import torch as t
from torch import nn
import numpy as np
class my_net(nn.Module):
    def __init__(self):
        super(my_net, self).__init__()
        self.sig1 = nn.Sigmoid()
        self.fc1 = nn.Linear(3, 2)
        self.fc2 = nn.Linear(2, 2)
        self.relu1 = nn.ReLU()

    def forward(self, x):
        x = self.sig1(x)
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.relu1(x)
        return x
def get_mask(data, threshold):
    mask = data > threshold
    data = data * mask
    return data
if __name__ == '__main__':
    inputs = t.tensor([1., 2., 3.], dtype=t.float)
    labels = t.tensor([0, 1], dtype=t.float)
    my_net1 = my_net()
    criterion = nn.CrossEntropyLoss()
    outputs = my_net1(inputs)
    loss = criterion(outputs, labels)
    weight_list = []
    bias_list = []
    for para in my_net1.parameters():
        if len(para.size()) == 2:
            weight_list.append(para)
        else:
            bias_list.append(para)
    print('weight_list', weight_list)

    for layer in my_net1.modules():
        for name, param in layer.named_parameters():
            if name == 'weight':
                layer.weight.data = get_mask(param, 1e-5)



    grads = t.autograd.grad(loss, my_net1.parameters(), retain_graph=True, create_graph=True)
    weight_list = []
    bias_list = []
    for para in my_net1.parameters():
        if len(para.size()) == 2:
            weight_list.append(para)
        else:
            bias_list.append(para)
    print('weight_list', weight_list)
    print('bias_list', bias_list)
    print('grads', grads)
    hessian_grads = []

    for k in range(len(grads)):
        hessian_grad = t.zeros_like(grads[k])
        for i in range(grads[k].size(0)):
            if len(grads[k].size()) == 2:
                for j in range(grads[k].size(1)):
                    hessian_grad[i, j] = t.autograd.grad(grads[k][i][j], my_net1.parameters(), retain_graph=True)[k][i, j]
            else:
                hessian_grad[i] = t.autograd.grad(grads[k][i], my_net1.parameters(), retain_graph=True)[k][i]
        hessian_grads.append(hessian_grad)

    print('hessian_grads', hessian_grads)
    weight_grad_total = []
    bias_grad_total = []
    for hessian_grad in hessian_grads:
        if len(hessian_grad.size()) == 2:
            weight_grad_total.append(hessian_grad)
        else:
            bias_grad_total.append(hessian_grad)

    print('weight_grad_total', weight_grad_total)
    print('bias_grad_total', bias_grad_total)
    fii = []
    for i in range(len(weight_list)):
        mask = weight_list[i] > 1e-5
        weight_list[i] = weight_list[i]*mask
        result = 1 / 2 * t.sqrt(weight_list[i]) * weight_grad_total[i]
        fii.append(result)
    print('fii', fii)




